package rabbit.flt.plugins.zookeeper.server;

import org.apache.zookeeper.Version;
import org.apache.zookeeper.server.ZooKeeperServer;
import rabbit.flt.common.AgentConfig;
import rabbit.flt.common.ScheduleTask;
import rabbit.flt.common.metrics.ZookeeperMetrics;

import java.lang.management.ManagementFactory;
import java.lang.management.MemoryMXBean;
import java.util.concurrent.ArrayBlockingQueue;

/**
 * zookeeper metrics 采集任务
 */
public class ZooKeeperMetricsTask extends ScheduleTask<ZookeeperMetrics> {

    private Long lastSampleTime = 0L;

    /**
     * zk server
     */
    private static ZooKeeperServer server;

    /**
     * 回调
     */
    private static ArrayBlockingQueue<Runnable> callBacks = new ArrayBlockingQueue<>(1024);

    @Override
    public boolean isPrepared(AgentConfig config) {
        if (null == server || null == server.getZKDatabase() || !server.isRunning()) {
            return false;
        }
        long interval = config.getZkMetricsSampleIntervalSeconds() * 1000L;
        long sampleTime = System.currentTimeMillis() / interval * interval;
        if (lastSampleTime != sampleTime) {
            lastSampleTime = sampleTime;
            return true;
        }
        return false;
    }

    @Override
    public boolean isMetricsEnabled(AgentConfig config) {
        return config.isZookeeperMetricsEnabled();
    }

    @Override
    public ZookeeperMetrics getMetrics() {
        ZookeeperMetrics metrics = loadMetrics();
        metrics.setSamplingTime(lastSampleTime);
        return metrics;
    }

    /**
     * zookeeper server启动了
     *
     * @param server
     */
    public static void serverStarted(ZooKeeperServer server) {
        ZooKeeperMetricsTask.server = server;
        Runnable job;
        while (null != (job = callBacks.poll())) {
            job.run();
        }
    }

    /**
     * 获取实时指标数据
     *
     * @return
     */
    public static ZookeeperMetrics loadMetrics() {
        ZookeeperMetrics metrics = new ZookeeperMetrics(server.getZKDatabase().getNodeCount(),
                server.getZKDatabase().getSessionCount(),
                server.getServerCnxnFactory().getNumAliveConnections(),
                server.getZKDatabase().getDataTree().getWatchCount());
        metrics.setState(server.serverStats().getServerState());
        metrics.setPendingRequests(server.getInflight());
        metrics.setMaxLatency(server.serverStats().getMaxLatency());
        metrics.setMinLatency(server.serverStats().getMinLatency());
        metrics.setAverageLatency((long) Math.floor(server.serverStats().getAvgLatency()));
        metrics.setThreadCount(ManagementFactory.getThreadMXBean().getThreadCount());
        MemoryMXBean mxBean = ManagementFactory.getMemoryMXBean();
        metrics.setMaxHeapMemory(mxBean.getHeapMemoryUsage().getMax() / MILLION);
        metrics.setUsedHeapMemory(mxBean.getHeapMemoryUsage().getUsed() / MILLION);
        metrics.setVersion(Version.getVersion());
        return metrics;
    }

    /**
     * 启动回调
     *
     * @param callback
     */
    public static void addStartupCallback(Runnable callback) {
        callBacks.add(callback);
    }
}
